# -*-coding:utf-8-*-
from PyQt5.Qt import *
import numpy as np
import cv2


class SegMode:
    def __init__(self, img, label_rect=None, pixmap_rect=None, scale_ratio=None):
        self.image = cv2.imread(img)
        self.label_rect = np.array(label_rect.getCoords())    # 获取label尺寸和鼠标位置
        self.pixmap_rect = np.array(pixmap_rect.getCoords())
        self.mouse_rect_coords = None
        self.scale_ratio = scale_ratio
        self.pixmap_loc = None
        self.img_loc = None
        self.processed_data = [self.image]

    # 求出框相对于原图位置
    def coord_transform(self, mouse_rect):
        margin_x = (self.label_rect[2] - self.pixmap_rect[2]) / 2
        margin_y = (self.label_rect[3] - self.label_rect[1] - self.pixmap_rect[3]) / 2
        margins = np.array((margin_x, margin_y, margin_x, margin_y))
        pixmap_loc = (mouse_rect - margins)   # 位置坐标从0开始，直接就是图像中的x, y index
        xl, xr = [np.clip(x, 0, self.pixmap_rect[2]) for x in (pixmap_loc[0], pixmap_loc[2])]
        yl, yr = [np.clip(y, 0, self.pixmap_rect[3]) for y in (pixmap_loc[1], pixmap_loc[3])]
        self.pixmap_loc = np.array((xl, yl, xr, yr), dtype=float)
        self.img_loc = (self.pixmap_loc / self.scale_ratio).astype(int)

    # 对局部区域做分割操作，并保留数据以粗糙度评估
    def paint_area(self, mouse_rect):
        mouse_rect = np.array(mouse_rect)
        self.coord_transform(mouse_rect)
        img = self.processed_data[-1]
        img[self.img_loc[1]:self.img_loc[3], self.img_loc[0]:self.img_loc[2], :] = 0
        self.processed_data.append(img)
        img = self.image_adaptation(self.processed_data[-1])
        qImage = self.return_qimage(img)
        return qImage

    def image_adaptation(self, img):
        ih, iw, _ = img.shape
        w, h = (self.label_rect[2] - self.label_rect[0] + 1),\
               (self.label_rect[3] - self.label_rect[1] + 1)
        # 保持纵横比
        if (ih / iw) >= (h / w):
            nw = int(iw / ih * h)
            new_image = cv2.resize(img, (nw, h))
        else:
            nh = int(ih / iw * w)
            new_image = cv2.resize(img, (w, nh))
        return new_image

    def select_area(self, mouse_rect):
        mouse_rect = np.array(mouse_rect)
        self.coord_transform(mouse_rect)
        img = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
        local_img = self.image[self.img_loc[1]:self.img_loc[3], self.img_loc[0]:self.img_loc[2], :]
        frame_img = cv2.rectangle(img, (self.img_loc[0], self.img_loc[1]), (self.img_loc[2], self.img_loc[3]),
                                  (255, 255, 255), 2)
        frame_img = self.image_adaptation(frame_img)
        qImage = QImage(frame_img.data, frame_img.shape[1], frame_img.shape[0],
                        frame_img.shape[2] * frame_img.shape[1], QImage.Format_RGB888)
        return qImage, local_img

    @staticmethod
    def return_qimage(img):
        image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        qImage = QImage(image.data, image.shape[1], image.shape[0],
                        image.shape[2] * image.shape[1], QImage.Format_RGB888)
        return qImage
    # TODO后续继续实现delete键删除框操作

    SEGMENT = 0
    FINE_SEGMENT = 1
    PAINT_AREA = 2
    SELECT_AREA = 3

